import torch
import torch.nn as nn

"""
输入层 ---->  隐藏层  ----->  输出层
多层线性层   ---->  多层感知机
"""


class MLP(nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        # 线性层
        self.fc1 = nn.Linear(input_size, 64)
        self.fc2 = nn.Linear(64, 128)
        self.fc3 = nn.Linear(128, output_size)
        # 激活函数
        self.relu = nn.ReLU()

    # 前向传播
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        out = self.fc3(x)
        return out


if __name__ == '__main__':
    model = MLP(input_size=4, output_size=2)
    x = torch.rand(10, 4)
    out = model(x)
    print(out)
    print(out.shape)

    # --------------------------------
    # 简化写法
    net = nn.Sequential(
        nn.Linear(4, 64),
        nn.ReLU(),
        nn.Linear(64, 128),
        nn.ReLU(),
        nn.Linear(128, 2)
    )
    out = net(x)
    print(out)
    print(out.shape)
